{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from model import *\n",
    "from reader import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "opts = Options()\n",
    "config = tf.ConfigProto()\n",
    "config.gpu_options.allow_growth = True\n",
    "\n",
    "sess = tf.InteractiveSession(config=config)\n",
    "\n",
    "#relation enhancement method\n",
    "use_jeval = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#parameters\n",
    "opts.hidden_size = 512\n",
    "opts.num_samples = 2048*3\n",
    "opts.keep_prob = 0.5\n",
    "opts.num_layers = 2\n",
    "opts.learning_rate=0.001"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### You have to select one of the three datasets,\n",
    "#### the default dataset is FB15K-237 (datasets\\[0\\])\n",
    "#### you can also select another two datasets.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#select one dataset\n",
    "datasets = ['237', 'FB15K', 'WN18']\n",
    "used_dataset = datasets[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if used_dataset == '237':\n",
    "    #set log filename and data path\n",
    "    file_name = '237-dskg-hs512'\n",
    "    opts.data_path = 'data/FB15k-237/'\n",
    "    \n",
    "    #different datasets use different data parser\n",
    "    model = FBRespective(opts, sess)\n",
    "if used_dataset == 'FB15K':\n",
    "    file_name = 'fb-dskg-hs512'\n",
    "    opts.data_path = 'data/FB15k-237/'\n",
    "    model = FBRespective(opts, sess)\n",
    "if used_dataset == 'WN18':\n",
    "    file_name = 'wn-dskg-hs512'\n",
    "    opts.data_path = 'data/wordnet-mlj12/'\n",
    "    model = WNRespective(opts, sess)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#calculate ranks\n",
    "def cal_ranks(probs, method, label):\n",
    "    if method == 'min':\n",
    "        probs = probs - probs[range(len(label)), label].reshape(len(probs), 1)\n",
    "        ranks = (probs > 0).sum(axis=1) + 1\n",
    "    else:\n",
    "        ranks = pd.DataFrame(probs).rank(axis=1, ascending=False, method=method)\n",
    "        ranks = ranks.values[range(len(label)), label]\n",
    "    return ranks\n",
    "\n",
    "#calculate performance\n",
    "def cal_performance(ranks, top=10):\n",
    "    m_r = sum(ranks) * 1.0 / len(ranks)\n",
    "    h_10 = sum(ranks <= top) * 1.0 / len(ranks)\n",
    "    mrr = (1. / ranks).sum() / len(ranks)\n",
    "    return m_r, h_10, mrr\n",
    "\n",
    "def eval_entity_prediction(model, data, filter_mat, method='min', return_ranks=False, return_probs=False, return_label_probs=False):\n",
    "    options = model._options\n",
    "    batch_size = options.batch_size\n",
    "    \n",
    "    label = data[:, 2]\n",
    "    \n",
    "    data, padding_num = model.padding_data(data)\n",
    "\n",
    "    num_batch = len(data) // batch_size \n",
    "    \n",
    "    e_placeholder, r_placeholder, fectch_entity_probs = model._eval_e, model._eval_r, model._entity_probs\n",
    "    \n",
    "    probs = []\n",
    "    for i in range(num_batch):\n",
    "        e = data[:, 0][i * batch_size:(i + 1) * batch_size]\n",
    "        r = data[:, 1][i * batch_size:(i + 1) * batch_size]\n",
    "        \n",
    "        feed_dict = {}\n",
    "        feed_dict[e_placeholder] = e\n",
    "        feed_dict[r_placeholder] = r\n",
    "        \n",
    "        probs.append(sess.run(fectch_entity_probs, feed_dict))\n",
    "    probs = np.concatenate(probs)[:len(data) - padding_num]\n",
    "\n",
    "    if return_label_probs:\n",
    "        return probs[range(len(label)), label]\n",
    "    \n",
    "    if return_probs:\n",
    "        return probs\n",
    "\n",
    "    filter_probs = probs * filter_mat\n",
    "    filter_probs[range(len(label)), label] = probs[range(len(label)), label]\n",
    "\n",
    "    filter_ranks = cal_ranks(filter_probs, method=method, label=label)\n",
    "    if return_ranks:\n",
    "        return filter_ranks\n",
    "    ranks = cal_ranks(probs, method=method, label=label)\n",
    "    m_r, h_10, mrr = cal_performance(ranks)\n",
    "    f_m_r, f_h_10, f_mrr = cal_performance(filter_ranks)\n",
    "    \n",
    "    return (m_r, h_10, mrr, f_m_r, f_h_10, f_mrr)\n",
    "\n",
    "def eval_relation_prediction(model, data, filter_mat, method='min', return_ranks=False, return_probs=False):\n",
    "    options = model._options\n",
    "    batch_size = options.batch_size\n",
    "    \n",
    "    #data[:, 0]-->e, data[:, 1]-->r, data[:, 2]-->e2\n",
    "    label = data[:, 1]\n",
    "    \n",
    "    data, padding_num = model.padding_data(data)\n",
    "\n",
    "    num_batch = len(data) // batch_size\n",
    "    \n",
    "    e_placeholder, fectch_relation_probs = model._eval_e, model._relation_probs\n",
    "    \n",
    "    probs = []\n",
    "    \n",
    "    for i in range(num_batch):\n",
    "        e = data[:, 0][i * batch_size:(i + 1) * batch_size]\n",
    "        \n",
    "        feed_dict = {}\n",
    "        feed_dict[e_placeholder] = e\n",
    "        \n",
    "        probs.append(sess.run(fectch_relation_probs, feed_dict))\n",
    "        \n",
    "    probs = np.concatenate(probs)[:len(data) - padding_num]\n",
    "    return probs\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#preprocess data\n",
    "\n",
    "test_data = np.array(model._test_data[['h_id', 'r_id', 't_id']].values)\n",
    "train_data = model._train_data[['h_id', 'r_id', 't_id']].values\n",
    "valid_data = model._valid_data[['h_id', 'r_id', 't_id']].values\n",
    "\n",
    "filter_mat = model._tail_test_filter_mat\n",
    "vfilter_mat = model._tail_valid_filter_mat\n",
    "\n",
    "all_data = np.concatenate([train_data, test_data,valid_data])\n",
    "p_data = np.concatenate([test_data,valid_data])\n",
    "\n",
    "def gen_rev_rel(test_data):\n",
    "    half = len(test_data)//2\n",
    "    forward = test_data[:half]\n",
    "    back = test_data[half:]\n",
    "    rev_rel_test_data = test_data[:]\n",
    "    rev_rel = np.concatenate([back[:,1], forward[:,1]])\n",
    "    return rev_rel\n",
    "\n",
    "rev_rel = gen_rev_rel(test_data)\n",
    "vrev_rel=  gen_rev_rel(valid_data)\n",
    "\n",
    "rev_rel_test_data = np.stack([np.arange(model._entity_num),np.arange(model._entity_num)], axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def cal_r(probs, label, filter_mat):\n",
    "    filter_probs = probs * filter_mat\n",
    "    \n",
    "    filter_probs[range(len(label)), label] = probs[range(len(label)), label]\n",
    "    filter_ranks = cal_ranks(filter_probs, method='min', label=label)\n",
    "    \n",
    "    return filter_ranks\n",
    "\n",
    "\n",
    "def joint_eval(test_data, filter_mat, rev_rel):\n",
    "    label=test_data[:, 2]\n",
    "\n",
    "    ep =  eval_entity_prediction(model, data=test_data, filter_mat=filter_mat, return_probs=True)\n",
    "    efr = cal_r(ep, label, filter_mat)\n",
    "    if use_jeval:\n",
    "        rp = eval_relation_prediction(model, rev_rel_test_data, filter_mat=None, return_probs=True).T\n",
    "        rp = rp**0.33\n",
    "        rp = rp[rev_rel]\n",
    "        joint_probs = ep * rp\n",
    "        joint_fr = cal_r(joint_probs, label, filter_mat)\n",
    "    else:\n",
    "        joint_fr = efr\n",
    "    return joint_fr, efr\n",
    "\n",
    "def joint_eval_raw(test_data, filter_mat, rev_rel):\n",
    "    label=test_data[:, 2]\n",
    "    \n",
    "    \n",
    "    ep =  eval_entity_prediction(model, data=test_data, filter_mat=filter_mat, return_probs=True)\n",
    "    efr = cal_ranks(ep, method='min', label=label)\n",
    "    if use_jeval:\n",
    "        rp = eval_relation_prediction(model, rev_rel_test_data, filter_mat=None, return_probs=True).T\n",
    "        rp = rp**0.33\n",
    "        rp = rp[rev_rel]\n",
    "        joint_probs = ep * rp\n",
    "        joint_fr = cal_ranks(joint_probs, method='min', label=label)\n",
    "    else:\n",
    "        joint_fr = efr\n",
    "    return joint_fr, efr\n",
    "\n",
    "def process_ranks(efr, i=0, last_mean_loss=1000, title=''):\n",
    "\n",
    "    MR, H1, MRR = cal_performance(efr[:len(efr)], top=1)\n",
    "    _, H10, _ = cal_performance(efr[:len(efr)], top=10)\n",
    "    msg = '%s epoch:%i, Hits@1:%.3f, Hits@10:%.3f, MR:%.3f, MRR:%.3f, mean_loss:%.3f' % (format(title,'<15'), i, H1, H10, MR, MRR, last_mean_loss)\n",
    "    print(msg)\n",
    "    return (i, H1, H10, MR, MRR, last_mean_loss)\n",
    "\n",
    "def handle_eval(i=0, last_mean_loss=1000, valid=True, test=False):\n",
    "    if valid:\n",
    "        jfr, efr = joint_eval(test_data=valid_data, filter_mat=vfilter_mat, rev_rel=vrev_rel)\n",
    "        jrr, rr = joint_eval_raw(test_data=valid_data, filter_mat=vfilter_mat, rev_rel=vrev_rel)\n",
    "        \n",
    "        process_ranks(rr, i, last_mean_loss, title='Valid-R')\n",
    "        process_ranks(jrr, i, last_mean_loss, title='Valid-R-RH')\n",
    "        \n",
    "        msg = process_ranks(efr, i, last_mean_loss, title='Valid-F')\n",
    "        jmsg = process_ranks(jfr, i, last_mean_loss, title='Valid-F-RH')\n",
    "        \n",
    "        #process early stop\n",
    "        current_hits_1 = jmsg[1]\n",
    "        if current_hits_1 > best_hits_1:\n",
    "            best_hits_1 = current_hits_1\n",
    "            dropped_time = 0\n",
    "        else:\n",
    "            dropped_time += 1\n",
    "        \n",
    "        \n",
    "        valid_results.append(msg)\n",
    "        valid_results.append(jmsg)\n",
    "        if i % 50 == 0:\n",
    "            pd.DataFrame(valid_results, columns=['epoch','Hits@1', 'Hits@10', 'MR', 'MRR', 'mean_loss']).to_csv('results/'+file_name+'valid')\n",
    "        \n",
    "    if test:\n",
    "        jfr, efr = joint_eval(test_data=test_data, filter_mat=filter_mat, rev_rel=rev_rel)\n",
    "        jrr, rr = joint_eval_raw(test_data=test_data, filter_mat=filter_mat, rev_rel=rev_rel)\n",
    "        \n",
    "        process_ranks(rr, i, last_mean_loss, title='Test-R')\n",
    "        process_ranks(jrr, i, last_mean_loss, title='Test-R-RH')\n",
    "        \n",
    "        \n",
    "        msg = process_ranks(efr, i, last_mean_loss, title='Test-F')\n",
    "        jmsg = process_ranks(jfr, i, last_mean_loss, title='Test-F-RH')\n",
    "        results.append(msg)\n",
    "        results.append(jmsg)\n",
    "        if i % 50 == 0:\n",
    "            pd.DataFrame(results, columns=['epoch','Hits@1', 'Hits@10', 'MR', 'MRR', 'mean_loss']).to_csv('results/'+file_name+'test')\n",
    "    return "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "epoch =0\n",
    "results = []\n",
    "valid_results = []\n",
    "last_mean_loss=1000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#early stop setting\n",
    "best_hits_1 = 0\n",
    "dropped_time = 0\n",
    "max_dropped_time = 3\n",
    "\n",
    "max_epoch = 300"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## the function handle_eval(i=i, last_mean_loss=last_mean_loss, valid=True, test=True) returns the evaluation results:\n",
    "\n",
    "**Valid** and **Test** denote the datasets\n",
    "\n",
    "**R** denotes Raw results\n",
    "\n",
    "**F** denotes Filtered results\n",
    "\n",
    "**RH** denotes using relation enhancement method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "\n",
    "for i in range(epoch, max_epoch):\n",
    "    if i % 20 == 0:\n",
    "        handle_eval(i=i, last_mean_loss=last_mean_loss, valid=True, test=False)\n",
    "    last_mean_loss = model.train()\n",
    "    epoch += 1\n",
    "    \n",
    "    #early stop\n",
    "    if dropped_time >= max_dropped_time:\n",
    "        break\n",
    "    \n",
    "handle_eval(i=i, last_mean_loss=last_mean_loss, valid=True, test=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
